#!/bin/bash

# 进入脚本所在目录
cd "$(dirname "$0")"

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 函数：显示菜单
show_menu() {
    clear
    echo -e "${BLUE}"
    echo "=========================================="
    echo "    股票数据采集系统 - 主菜单"
    echo "=========================================="
    echo -e "${NC}"
    echo "1. 初始化数据库"
    echo "2. 更新股票基本信息"
    echo "3. 更新日线数据（所有股票，最近1天）"
    echo "4. 更新日线数据（所有股票，最近5天）"
    echo "5. 更新指定股票数据"
    echo "6. 查看数据库状态"
    echo "7. 运行维护任务"
    echo "8. 退出"
    echo ""
    echo -e "${YELLOW}请选择操作 [1-8]: ${NC}"
}

# 函数：等待用户按键
press_any_key() {
    echo ""
    echo -e "${YELLOW}按任意键继续...${NC}"
    read -n 1 -s -r
}

# 函数：运行命令并等待
run_command() {
    local command=$1
    local description=$2

    echo -e "${BLUE}执行: $description${NC}"
    echo "------------------------------------------"

    # 执行命令
    $command

    local exit_code=$?
    echo "------------------------------------------"

    if [ $exit_code -eq 0 ]; then
        echo -e "${GREEN}✅ 执行成功${NC}"
    elif [ $exit_code -eq 130 ]; then
        echo -e "${YELLOW}⚠️  操作被用户中断${NC}"
    else
        echo -e "${RED}❌ 执行失败 (退出代码: $exit_code)${NC}"
    fi

    press_any_key
}

# 主循环
while true; do
    show_menu
    read choice

    case $choice in
        1)
            run_command "python3 ../main.py --init" "初始化数据库"
            ;;
        2)
            run_command "python3 ../main.py --update-basic" "更新股票基本信息"
            ;;
        3)
            run_command "python3 ../main.py --update-daily" "更新日线数据（最近1天）"
            ;;
        4)
            run_command "python3 ../main.py --update-daily 5" "更新日线数据（最近5天）"
            ;;
        5)
            echo -e "${YELLOW}请输入股票代码: ${NC}"
            read symbol
            if [ -n "$symbol" ]; then
                run_command "python3 ../main.py --update-daily --symbol $symbol" "更新股票 $symbol 的数据"
            else
                echo -e "${RED}❌ 股票代码不能为空${NC}"
                press_any_key
            fi
            ;;
        6)
            run_command "python3 ../main.py --status" "查看数据库状态"
            ;;
        7)
            run_command "python3 ../main.py --maintenance" "运行维护任务"
            ;;
        8)
            echo -e "${GREEN}感谢使用，再见！${NC}"
            exit 0
            ;;
        *)
            echo -e "${RED}❌ 无效选择，请重新输入${NC}"
            press_any_key
            ;;
    esac
done